#!/usr/bin/env python3

"""
MCP HTTP Client - Python implementation following MCP HTTP Transport specification
This client communicates with MCP servers using JSON-RPC 2.0 over HTTP
"""

import requests
import json
import sys
from typing import Optional, Dict, Any, List

class MCPHttpClient:
    def __init__(self, server_url: str = None):
        if server_url is None:
            import os
            server_url = os.getenv('MCP_SERVER_URL', 'http://localhost:3000/mcp')
        self.server_url = server_url
        self.session_id = None
        self.request_id = 1
        
        # Configure session
        self.session = requests.Session()
        self.session.headers.update({
            'Content-Type': 'application/json',
            'Accept': 'application/json, text/event-stream'
        })
        self.session.timeout = 60

    def get_next_request_id(self) -> str:
        """Generate next request ID"""
        request_id = str(self.request_id)
        self.request_id += 1
        return request_id

    def send_request(self, method: str, params: Dict[str, Any] = None) -> Any:
        """Send JSON-RPC request to MCP server"""
        if params is None:
            params = {}
            
        request = {
            "jsonrpc": "2.0",
            "id": self.get_next_request_id(),
            "method": method,
            "params": params
        }

        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json, text/event-stream'
        }

        # Add session ID if available
        if self.session_id:
            headers['Mcp-Session-Id'] = self.session_id

        try:
            response = self.session.post(self.server_url, json=request, headers=headers)
            response.raise_for_status()
            
            # Validate JSON-RPC response
            data = response.json()
            if not data.get('jsonrpc') or data['jsonrpc'] != '2.0':
                raise Exception('Invalid JSON-RPC response format')

            if 'error' in data:
                error = data['error']
                raise Exception(f"MCP Error {error['code']}: {error['message']}")

            return data.get('result')
        
        except requests.RequestException as e:
            if hasattr(e, 'response') and e.response is not None:
                print(f"HTTP Error: {e.response.status_code} {e.response.reason}")
                try:
                    error_data = e.response.json()
                    print(f"Response: {json.dumps(error_data, indent=2)}")
                except:
                    print(f"Response: {e.response.text}")
            raise

    def initialize(self) -> Dict[str, Any]:
        """Initialize connection with MCP server"""
        print("🔌 Initializing MCP connection...")
        
        result = self.send_request('initialize', {
            'protocolVersion': '2025-03-26',
            'capabilities': {
                'tools': True,
                'resources': False,
                'prompts': False
            },
            'clientInfo': {
                'name': 'mcp-http-client-python',
                'version': '1.0.0'
            }
        })

        self.session_id = result.get('sessionId')
        print("✅ MCP connection established")
        print(f"   Protocol Version: {result['protocolVersion']}")
        print(f"   Server: {result['serverInfo']['name']} v{result['serverInfo']['version']}")
        if self.session_id:
            print(f"   Session ID: {self.session_id}")

        return result

    def list_tools(self) -> Dict[str, Any]:
        """List available tools"""
        return self.send_request('tools/list')

    def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """Call a specific tool"""
        return self.send_request('tools/call', {
            'name': name,
            'arguments': arguments
        })

    def disconnect(self):
        """Disconnect from server"""
        if self.session_id:
            try:
                base_url = self.server_url.replace('/mcp', '')
                self.session.delete(f"{base_url}/mcp/session/{self.session_id}")
                print("🔌 Session terminated")
            except Exception as e:
                print(f"Warning: Failed to terminate session: {e}")
            self.session_id = None

    # Convenience methods for Java Decompiler tools

    def search_maven_artifacts(self, query: str) -> Dict[str, Any]:
        return self.call_tool('search_maven_artifacts', {'query': query})

    def download_artifact(self, group_id: str, artifact_id: str, version: str, packaging: str = 'jar') -> Dict[str, Any]:
        return self.call_tool('download_artifact', {
            'groupId': group_id,
            'artifactId': artifact_id,
            'version': version,
            'packaging': packaging
        })

    def search_classes(self, query: str, artifact_id: Optional[str] = None) -> Dict[str, Any]:
        args = {'query': query}
        if artifact_id:
            args['artifactId'] = artifact_id
        return self.call_tool('search_classes', args)

    def decompile_class(self, group_id: str, artifact_id: str, version: str, class_name: str) -> Dict[str, Any]:
        return self.call_tool('decompile_class', {
            'groupId': group_id,
            'artifactId': artifact_id,
            'version': version,
            'className': class_name
        })

    def get_artifact_versions(self, group_id: str, artifact_id: str) -> Dict[str, Any]:
        return self.call_tool('get_artifact_versions', {
            'groupId': group_id,
            'artifactId': artifact_id
        })

    def scan_local_project(self, project_path: str) -> Dict[str, Any]:
        return self.call_tool('scan_local_project', {
            'projectPath': project_path
        })

    def search_local_repository(self, query: str) -> Dict[str, Any]:
        return self.call_tool('search_local_maven_repository', {
            'query': query
        })

    def find_class_in_local_repository(self, class_name: str, repository_path: Optional[str] = None) -> Dict[str, Any]:
        args = {'className': class_name}
        if repository_path:
            args['repositoryPath'] = repository_path
        return self.call_tool('find_class_in_local_repository', args)


def run_mcp_example():
    """Example usage of MCP HTTP Client"""
    client = MCPHttpClient()

    try:
        print("🧪 Testing MCP HTTP Client...\n")

        # Initialize connection
        client.initialize()
        print()

        # List available tools
        print("📋 Listing available tools...")
        tools = client.list_tools()
        print(f"   Found {len(tools['tools'])} tools:")
        for tool in tools['tools']:
            print(f"   - {tool['name']}: {tool['description']}")
        print()

        # Search for Maven artifacts
        print("🔍 Searching for Maven artifacts (commons-lang3)...")
        search_result = client.search_maven_artifacts('commons-lang3')
        search_data = json.loads(search_result['content'][0]['text'])
        print(f"   Found {search_data['found']} artifacts")
        
        if search_data['artifacts']:
            first_artifact = search_data['artifacts'][0]
            print(f"   First artifact: {first_artifact['groupId']}:{first_artifact['artifactId']}:{first_artifact['version']}")
            
            # Download artifact
            print("\n📦 Downloading artifact...")
            download_result = client.download_artifact(
                first_artifact['groupId'],
                first_artifact['artifactId'],
                first_artifact['version']
            )
            
            download_data = json.loads(download_result['content'][0]['text'])
            if download_data['success']:
                print(f"   ✅ Downloaded successfully! Found {download_data['classCount']} classes")
                
                # Search for classes
                print("\n🔎 Searching for classes (String)...")
                class_search_result = client.search_classes('String')
                class_data = json.loads(class_search_result['content'][0]['text'])
                
                if class_data['totalMatches'] > 0:
                    print(f"   Found {class_data['totalMatches']} matching classes")
                    first_class = class_data['matches'][0]
                    print(f"   First match: {first_class['fullName']}")
                    
                    # Decompile class
                    print("\n🛠️  Decompiling class...")
                    decompile_result = client.decompile_class(
                        first_class['artifact']['groupId'],
                        first_class['artifact']['artifactId'],
                        first_class['artifact']['version'],
                        first_class['fullName']
                    )
                    
                    source_code = decompile_result['content'][0]['text']
                    if 'Decompilation failed' in source_code:
                        print("   ⚠️ Decompilation failed")
                    else:
                        print("   ✅ Decompilation successful!")
                        print(f"   📄 Source code length: {len(source_code)} characters")
                        print("   📝 First 300 characters:")
                        print(f"   {source_code[:300]}...")

        # Get artifact versions
        print("\n📌 Getting versions for commons-lang3...")
        version_result = client.get_artifact_versions('org.apache.commons', 'commons-lang3')
        version_data = json.loads(version_result['content'][0]['text'])
        print(f"   Found {len(version_data['versions'])} versions")
        print(f"   Latest versions: {', '.join(version_data['versions'][:5])}")

        # Search local repository
        print("\n🏠 Searching local repository...")
        local_result = client.search_local_repository('commons')
        local_data = json.loads(local_result['content'][0]['text'])
        print(f"   Found {local_data['found']} local artifacts")

        print("\n✅ MCP HTTP Client test completed successfully!")

    except Exception as error:
        print(f"\n❌ MCP HTTP Client test failed: {error}")
        
        if 'Connection refused' in str(error):
            print("\n💡 Make sure the MCP HTTP server is running:")
            print("   npm run build")
            print("   npm run start:mcp-http")
    
    finally:
        # Clean up
        client.disconnect()


def main():
    """Command line interface"""
    if len(sys.argv) > 1:
        client = MCPHttpClient()
        command = sys.argv[1]
        args = sys.argv[2:]

        try:
            client.initialize()

            if command == 'tools':
                tools = client.list_tools()
                print(json.dumps(tools, indent=2))

            elif command == 'search':
                if not args:
                    print("Usage: python mcp-http-client.py search <query>")
                    sys.exit(1)
                result = client.search_maven_artifacts(args[0])
                print(result['content'][0]['text'])

            elif command == 'download':
                if len(args) < 3:
                    print("Usage: python mcp-http-client.py download <groupId> <artifactId> <version>")
                    sys.exit(1)
                result = client.download_artifact(args[0], args[1], args[2])
                print(result['content'][0]['text'])

            elif command == 'decompile':
                if len(args) < 4:
                    print("Usage: python mcp-http-client.py decompile <groupId> <artifactId> <version> <className>")
                    sys.exit(1)
                result = client.decompile_class(args[0], args[1], args[2], args[3])
                print(result['content'][0]['text'])

            elif command == 'example':
                run_mcp_example()

            else:
                print("Available commands:")
                print("  tools                                           - List available tools")
                print("  search <query>                                  - Search Maven artifacts")
                print("  download <groupId> <artifactId> <version>       - Download artifact")
                print("  decompile <groupId> <artifactId> <version> <className> - Decompile class")
                print("  example                                         - Run example workflow")

        except Exception as error:
            print(f"Error: {error}")
            sys.exit(1)
        finally:
            client.disconnect()
    else:
        # Run example by default
        run_mcp_example()


if __name__ == "__main__":
    main()